#include "evdeploy/cv/detection/yolov8_prepostprocessor.h"

namespace ev
{
    namespace vision
    {
        float YOLOv8Preprocessor::GetScale() { return m_scale; }

        EVStatus YOLOv8Preprocessor::Run(cv::Mat &in_mat, cv::Mat &out_mat, int out_width, bool normalize, bool bgr2rgb)
        {
            REC_TIME(t0);
            if (out_width < 10)
            {
                EVLOG(ERROR) << "input size too small:" << out_width;
                return EV_FAIL;
            }
            m_scale =
                std::min(out_width / static_cast<float>(in_mat.rows), out_width / static_cast<float>(in_mat.cols));
            cv::Size new_size = cv::Size{in_mat.cols * m_scale, in_mat.rows * m_scale};
            cv::Mat resized_mat;
            cv::resize(in_mat, resized_mat, new_size);

            cv::Mat board = cv::Mat(cv::Size(out_width, out_width), CV_8UC3, cv::Scalar(114, 114, 114));

            resized_mat.copyTo(board(cv::Rect{0, 0, resized_mat.cols, resized_mat.rows}));

            if (bgr2rgb)
            {
                cv::cvtColor(board, board, cv::COLOR_BGR2RGB);
            }

            if (normalize)
            {
                board.convertTo(board, CV_32F, 1 / 255.);
            }
            else
            {
                board.convertTo(board, CV_32F);
            }

            out_mat = board.clone();
            std::vector<cv::Mat> chw_wrappers;
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1, out_mat.data);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1,
                                      out_mat.data + sizeof(float) * out_width * out_width);
            chw_wrappers.emplace_back(out_width, out_width, CV_32FC1,
                                      out_mat.data + 2 * sizeof(float) * out_width * out_width);
            cv::split(board, chw_wrappers);

            REC_TIME(t1);
            EVLOG(INFO) << "YOLOv8Preprocessor run time(ms):" << RUN_TIME(t1 - t0);

            return EV_SUCCESS;
        }

        EVStatus YOLOv8Postprocessor::Run(EVMatData *out, std::vector<ev::vision::BoxInfo> &objects, float scale,
                                          float thresh, const int img_w, const int img_h, float nms_thresh, bool class_aware)
        {

            REC_TIME(t0);
            int box_num = out->dims[2];
            int channel = out->dims[1];
            int class_num = channel - 4;
            float *data = (float *)out->data;
            for (int i = 0; i < box_num; ++i)
            {
                float max_score = data[4 * box_num + i];
                int index = 0;
                for (int j = 5; j < channel; j++)
                {
                    if (data[i + j * box_num] > max_score)
                    {
                        max_score = data[i + j * box_num];
                        index = j - 4;
                    }
                }
                float x = data[i];
                float y = data[box_num + i];
                float w = data[2 * box_num + i];
                float h = data[3 * box_num + i];
                if (max_score > thresh)
                {
                    float left = (x - 0.5 * w) / scale;
                    float top = (y - 0.5 * h) / scale;
                    float width = w / scale;
                    float height = h / scale;
                    BoxInfo box_info = {left, top, left + width, top + height, max_score, index};
                    box_info.x1 = std::max(box_info.x1, 0.0f);
                    box_info.x2 = std::min(box_info.x2, img_w - 1.0f);
                    box_info.y1 = std::max(box_info.y1, 0.0f);
                    box_info.y2 = std::min(box_info.y2, img_h - 1.0f);
                    objects.push_back(box_info);
                }
            }
            NMS(objects, nms_thresh, class_aware);
            REC_TIME(t1);
            EVLOG(INFO) << "YOLOv8Postprocessor run time(ms):" << RUN_TIME(t1 - t0);
            return EV_SUCCESS;
        }

    } // namespace vision
} // namespace ev